import sys
import os
import time
import threading
from datetime import datetime, timedelta
import mss
import numpy as np
import cv2
from PyQt5.QtWidgets import (
    QApplication, QMainWindow, QPushButton, QVBoxLayout, QHBoxLayout, 
    QWidget, QLabel, QFileDialog, QMessageBox, QComboBox, QSpinBox,
    QCheckBox, QGroupBox, QSlider, QSystemTrayIcon, QMenu, QAction,
    QStyle, QSizePolicy, QProgressBar
)
from PyQt5.QtCore import Qt, QThread, pyqtSignal, QTimer, QSize
from PyQt5.QtGui import QImage, QPixmap, QIcon, QKeySequence
from PyQt5.QtMultimedia import QAudioRecorder, QAudioEncoderSettings

# 检查依赖包
try:
    import pkg_resources
    required_packages = ['PyQt5', 'opencv-python', 'mss', 'numpy']
    installed_packages = {pkg.key for pkg in pkg_resources.working_set}
    missing_packages = [pkg for pkg in required_packages if pkg.lower() not in installed_packages]
    
    if missing_packages:
        print(f"缺少以下依赖包: {', '.join(missing_packages)}")
        print("请使用以下命令安装依赖:")
        print(f"pip install -i https://pypi.tuna.tsinghua.edu.cn/simple {' '.join(missing_packages)}")
        sys.exit(1)
except ImportError:
    pass  # 如果pkg_resources不可用，跳过检查

class RecordingTimer(QThread):
    """录制时间计时器线程"""
    update_time = pyqtSignal(int)  # 发送已录制的秒数
    
    def __init__(self):
        super().__init__()
        self.running = False
        self.seconds = 0
    
    def run(self):
        self.running = True
        self.seconds = 0
        while self.running:
            self.seconds += 1
            self.update_time.emit(self.seconds)
            time.sleep(1)
    
    def stop(self):
        self.running = False
        self.wait()

class ScreenRecorderThread(QThread):
    """屏幕录制线程类"""
    # 定义信号
    update_frame = pyqtSignal(np.ndarray)  # 用于更新预览帧
    recording_stopped = pyqtSignal()  # 用于通知录制已停止
    recording_error = pyqtSignal(str)  # 用于通知录制错误
    recording_progress = pyqtSignal(int)  # 用于更新录制进度

    def __init__(self, output_file, fps=15.0, preview=False, monitor_num=1, record_audio=False):
        super().__init__()
        self.output_file = output_file
        self.fps = fps
        self.preview = preview
        self.monitor_num = monitor_num - 1  # mss使用0为第一个显示器
        self.record_audio = record_audio
        self.running = False
        self.audio_recorder = None
        self.audio_file = None
        
        # 进度更新计时器
        self.progress_timer = None
        self.start_time = None

    def setup_audio_recorder(self):
        """设置音频录制"""
        if not self.record_audio:
            return
            
        try:
            self.audio_file = os.path.splitext(self.output_file)[0] + ".wav"
            self.audio_recorder = QAudioRecorder()
            settings = QAudioEncoderSettings()
            settings.setCodec("audio/pcm")
            settings.setQuality(QAudioEncoderSettings.HighQuality)
            self.audio_recorder.setEncodingSettings(settings)
            self.audio_recorder.setOutputLocation(self.audio_file)
        except Exception as e:
            self.recording_error.emit(f"音频录制设置失败: {str(e)}")
            self.record_audio = False

    def run(self):
        """线程运行函数"""
        try:
            self.running = True
            self.start_time = time.time()
            
            # 设置进度更新计时器
            self.progress_timer = QTimer()
            self.progress_timer.timeout.connect(self.update_progress)
            self.progress_timer.start(1000)  # 每秒更新一次
            
            # 设置音频录制
            if self.record_audio:
                self.setup_audio_recorder()
                if self.audio_recorder:
                    self.audio_recorder.record()
            
            # 获取屏幕大小
            with mss.mss() as sct:
                try:
                    monitor = sct.monitors[self.monitor_num]
                except IndexError:
                    self.recording_error.emit(f"显示器 {self.monitor_num + 1} 不存在，使用主显示器")
                    monitor = sct.monitors[0]  # 使用主显示器
                
                width = monitor['width']
                height = monitor['height']

            # 设置视频编码器和输出
            fourcc = cv2.VideoWriter_fourcc(*"mp4v")
            out = cv2.VideoWriter(self.output_file, fourcc, self.fps, (width, height))

            # 开始录制
            start_time = time.time()
            frame_count = 0

            with mss.mss() as sct:
                while self.running:
                    # 捕获屏幕
                    screenshot = sct.grab(monitor)
                    # 转换为numpy数组
                    frame = np.array(screenshot)
                    # 转换颜色格式 (BGRA to BGR)
                    frame = cv2.cvtColor(frame, cv2.COLOR_BGRA2BGR)
                    # 写入视频
                    out.write(frame)
                    frame_count += 1

                    # 发送信号更新预览
                    if self.preview and frame_count % 3 == 0:  # 每3帧更新一次预览，减少CPU负担
                        self.update_frame.emit(frame)

                    # 控制帧率
                    elapsed = time.time() - start_time
                    expected_frames = int(elapsed * self.fps)
                    if frame_count > expected_frames:
                        self.msleep(int((frame_count - expected_frames) * 1000 / self.fps))

            # 释放资源
            out.release()
            
            # 停止音频录制
            if self.record_audio and self.audio_recorder:
                self.audio_recorder.stop()
                
                # 合并音视频（如果有音频）
                if os.path.exists(self.audio_file) and os.path.getsize(self.audio_file) > 0:
                    self.merge_audio_video()
            
            if self.progress_timer:
                self.progress_timer.stop()
                
            self.recording_stopped.emit()
            
        except Exception as e:
            self.recording_error.emit(f"录制过程中出错: {str(e)}")
            if self.progress_timer:
                self.progress_timer.stop()
            self.recording_stopped.emit()

    def update_progress(self):
        """更新录制进度"""
        if self.start_time:
            elapsed = int(time.time() - self.start_time)
            self.recording_progress.emit(elapsed)

    def merge_audio_video(self):
        """合并音频和视频文件"""
        try:
            # 这里可以使用ffmpeg合并音视频，但需要安装ffmpeg
            # 为简化依赖，这里只保留音频文件，让用户自行合并
            pass
        except Exception as e:
            self.recording_error.emit(f"合并音视频失败: {str(e)}")

    def stop(self):
        """停止录制"""
        self.running = False
        if self.audio_recorder:
            self.audio_recorder.stop()
        self.wait()

class ScreenRecorderApp(QMainWindow):
    """录屏软件主窗口类"""
    def __init__(self):
        super().__init__()
        self.output_file = None
        self.recorder_thread = None
        self.timer_thread = None
        self.recording_duration = 0
        self.tray_icon = None
        self.init_ui()
        self.setup_tray_icon()
        self.setup_shortcuts()
        
        # 设置应用图标
        self.setWindowIcon(self.style().standardIcon(QStyle.SP_ComputerIcon))
        
        # 检测显示器
        self.detect_monitors()

    def init_ui(self):
        """初始化用户界面"""
        self.setWindowTitle("mars录屏")
        self.setGeometry(100, 100, 600, 500)
        self.setMinimumSize(500, 400)

        # 创建主布局
        central_widget = QWidget()
        main_layout = QVBoxLayout(central_widget)
        self.setCentralWidget(central_widget)

        # 添加标题标签
        title_label = QLabel("mars录屏")
        title_label.setAlignment(Qt.AlignCenter)
        title_label.setStyleSheet("font-size: 20pt; font-weight: bold; margin: 10px;")
        main_layout.addWidget(title_label)

        # 添加设置区域
        settings_group = QGroupBox("录制设置")
        settings_layout = QVBoxLayout(settings_group)
        
        # 添加文件路径设置
        path_layout = QHBoxLayout()
        path_layout.addWidget(QLabel("保存路径:"))
        self.path_label = QLabel("未选择保存路径")
        self.path_label.setStyleSheet("color: gray; font-style: italic;")
        self.path_label.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Preferred)
        path_layout.addWidget(self.path_label)
        self.select_path_btn = QPushButton("选择路径")
        self.select_path_btn.clicked.connect(self.select_output_path)
        path_layout.addWidget(self.select_path_btn)
        settings_layout.addLayout(path_layout)
        
        # 添加显示器选择
        monitor_layout = QHBoxLayout()
        monitor_layout.addWidget(QLabel("录制显示器:"))
        self.monitor_combo = QComboBox()
        self.monitor_combo.addItem("主显示器")  # 默认选项
        monitor_layout.addWidget(self.monitor_combo)
        settings_layout.addLayout(monitor_layout)
        
        # 添加帧率设置
        fps_layout = QHBoxLayout()
        fps_layout.addWidget(QLabel("帧率:"))
        self.fps_spin = QSpinBox()
        self.fps_spin.setRange(1, 60)
        self.fps_spin.setValue(15)
        self.fps_spin.setSuffix(" FPS")
        fps_layout.addWidget(self.fps_spin)
        
        # 添加音频录制选项
        self.audio_checkbox = QCheckBox("录制系统音频")
        fps_layout.addWidget(self.audio_checkbox)
        settings_layout.addLayout(fps_layout)
        
        main_layout.addWidget(settings_group)

        # 添加状态区域
        status_group = QGroupBox("录制状态")
        status_layout = QVBoxLayout(status_group)
        
        # 添加状态标签
        self.status_label = QLabel("就绪")
        self.status_label.setAlignment(Qt.AlignCenter)
        self.status_label.setStyleSheet("color: green; font-size: 14pt; font-weight: bold;")
        status_layout.addWidget(self.status_label)
        
        # 添加录制时间显示
        time_layout = QHBoxLayout()
        time_layout.addWidget(QLabel("录制时间:"))
        self.time_label = QLabel("00:00:00")
        self.time_label.setStyleSheet("font-size: 12pt; font-family: monospace;")
        time_layout.addWidget(self.time_label)
        status_layout.addLayout(time_layout)
        
        # 添加进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setTextVisible(False)
        self.progress_bar.setRange(0, 100)
        self.progress_bar.setValue(0)
        self.progress_bar.setStyleSheet("""
            QProgressBar {
                border: 1px solid grey;
                border-radius: 5px;
                text-align: center;
            }
            QProgressBar::chunk {
                background-color: #05B8CC;
                width: 10px;
                margin: 0.5px;
            }
        """)
        status_layout.addWidget(self.progress_bar)
        
        main_layout.addWidget(status_group)

        # 添加按钮布局
        btn_layout = QHBoxLayout()

        # 添加开始录制按钮
        self.start_btn = QPushButton("开始录制")
        self.start_btn.setStyleSheet("""
            QPushButton {
                background-color: #4CAF50;
                color: white;
                font-size: 14pt;
                padding: 10px;
                border-radius: 5px;
            }
            QPushButton:hover {
                background-color: #45a049;
            }
            QPushButton:disabled {
                background-color: #cccccc;
            }
        """)
        self.start_btn.clicked.connect(self.start_recording)
        btn_layout.addWidget(self.start_btn)

        # 添加停止录制按钮
        self.stop_btn = QPushButton("停止录制")
        self.stop_btn.setStyleSheet("""
            QPushButton {
                background-color: #f44336;
                color: white;
                font-size: 14pt;
                padding: 10px;
                border-radius: 5px;
            }
            QPushButton:hover {
                background-color: #d32f2f;
            }
            QPushButton:disabled {
                background-color: #cccccc;
            }
        """)
        self.stop_btn.clicked.connect(self.stop_recording)
        self.stop_btn.setEnabled(False)  # 初始时禁用
        btn_layout.addWidget(self.stop_btn)

        main_layout.addLayout(btn_layout)

        # 添加预览区域
        preview_group = QGroupBox("预览")
        preview_layout = QVBoxLayout(preview_group)
        self.preview_label = QLabel("预览区域")
        self.preview_label.setAlignment(Qt.AlignCenter)
        self.preview_label.setStyleSheet("border: 1px solid #cccccc; background-color: #f0f0f0;")
        self.preview_label.setMinimumHeight(150)
        preview_layout.addWidget(self.preview_label)
        main_layout.addWidget(preview_group)
        
        # 添加快捷键提示
        shortcut_label = QLabel("快捷键: F9 = 开始/停止录制, Esc = 退出")
        shortcut_label.setAlignment(Qt.AlignCenter)
        shortcut_label.setStyleSheet("color: gray; font-style: italic;")
        main_layout.addWidget(shortcut_label)

    def setup_shortcuts(self):
        """设置键盘快捷键"""
        # F9 开始/停止录制
        self.record_shortcut = QKeySequence("F9")
        self.record_action = QAction("开始/停止录制", self)
        self.record_action.setShortcut(self.record_shortcut)
        self.record_action.triggered.connect(self.toggle_recording)
        self.addAction(self.record_action)
        
        # Esc 退出应用
        self.exit_shortcut = QKeySequence("Esc")
        self.exit_action = QAction("退出", self)
        self.exit_action.setShortcut(self.exit_shortcut)
        self.exit_action.triggered.connect(self.close)
        self.addAction(self.exit_action)

    def setup_tray_icon(self):
        """设置系统托盘图标"""
        try:
            self.tray_icon = QSystemTrayIcon(self)
            self.tray_icon.setIcon(self.style().standardIcon(QStyle.SP_ComputerIcon))
            
            # 创建托盘菜单
            tray_menu = QMenu()
            
            # 添加显示/隐藏动作
            show_action = QAction("显示", self)
            show_action.triggered.connect(self.show)
            tray_menu.addAction(show_action)
            
            # 添加开始/停止录制动作
            self.tray_record_action = QAction("开始录制", self)
            self.tray_record_action.triggered.connect(self.toggle_recording)
            tray_menu.addAction(self.tray_record_action)
            
            # 添加退出动作
            quit_action = QAction("退出", self)
            quit_action.triggered.connect(QApplication.quit)
            tray_menu.addAction(quit_action)
            
            # 设置托盘菜单
            self.tray_icon.setContextMenu(tray_menu)
            self.tray_icon.show()
            
            # 连接托盘图标点击事件
            self.tray_icon.activated.connect(self.tray_icon_activated)
        except Exception as e:
            print(f"系统托盘初始化失败: {str(e)}")

    def tray_icon_activated(self, reason):
        """托盘图标被激活时的处理"""
        if reason == QSystemTrayIcon.DoubleClick:
            if self.isVisible():
                self.hide()
            else:
                self.show()
                self.activateWindow()

    def detect_monitors(self):
        """检测系统显示器"""
        try:
            with mss.mss() as sct:
                self.monitor_combo.clear()
                for i, monitor in enumerate(sct.monitors[1:], 1):  # 跳过第一个（全部显示器）
                    width, height = monitor["width"], monitor["height"]
                    self.monitor_combo.addItem(f"显示器 {i} ({width}x{height})", i)
        except Exception as e:
            print(f"检测显示器失败: {str(e)}")
            self.monitor_combo.clear()
            self.monitor_combo.addItem("主显示器", 1)

    def select_output_path(self):
        """选择视频保存路径"""
        # 生成默认文件名 (基于当前时间)
        default_filename = datetime.now().strftime("录屏_%Y%m%d_%H%M%S.mp4")
        file_path, _ = QFileDialog.getSaveFileName(
            self, "保存视频", os.path.join(os.getcwd(), default_filename), "MP4 Files (*.mp4);;All Files (*)"
        )
        if file_path:
            self.output_file = file_path
            self.path_label.setText(file_path)
            self.path_label.setStyleSheet("color: blue;")

    def toggle_recording(self):
        """切换录制状态"""
        if self.recorder_thread and self.recorder_thread.isRunning():
            self.stop_recording()
        else:
            self.start_recording()

    def start_recording(self):
        """开始录制"""
        if not self.output_file:
            QMessageBox.warning(self, "警告", "请先选择保存路径")
            return

        # 检查文件是否已存在
        if os.path.exists(self.output_file):
            reply = QMessageBox.question(
                self, "确认", "文件已存在，是否覆盖？", QMessageBox.Yes | QMessageBox.No
            )
            if reply == QMessageBox.No:
                return

        # 获取设置
        fps = self.fps_spin.value()
        monitor_idx = self.monitor_combo.currentData()
        record_audio = self.audio_checkbox.isChecked()
        
        # 创建录制线程
        self.recorder_thread = ScreenRecorderThread(
            self.output_file, 
            fps=fps, 
            preview=True, 
            monitor_num=monitor_idx,
            record_audio=record_audio
        )
        
        # 连接信号
        self.recorder_thread.recording_stopped.connect(self.on_recording_stopped)
        self.recorder_thread.update_frame.connect(self.update_preview)
        self.recorder_thread.recording_error.connect(self.on_recording_error)
        self.recorder_thread.recording_progress.connect(self.update_progress_bar)
        
        # 创建计时器线程
        self.timer_thread = RecordingTimer()
        self.timer_thread.update_time.connect(self.update_recording_time)
        
        # 开始录制和计时
        self.recorder_thread.start()
        self.timer_thread.start()
        
        # 更新UI状态
        self.status_label.setText("正在录制...")
        self.status_label.setStyleSheet("color: red; font-size: 14pt; font-weight: bold;")
        self.start_btn.setEnabled(False)
        self.stop_btn.setEnabled(True)
        self.select_path_btn.setEnabled(False)
        self.fps_spin.setEnabled(False)
        self.monitor_combo.setEnabled(False)
        self.audio_checkbox.setEnabled(False)
        
        # 更新托盘图标
        if self.tray_icon:
            self.tray_icon.setToolTip("mars录屏 - 正在录制")
            self.tray_record_action.setText("停止录制")
            
        # 显示通知
        if self.tray_icon:
            self.tray_icon.showMessage("录制开始", "屏幕录制已开始", QSystemTrayIcon.Information, 2000)

    def stop_recording(self):
        """停止录制"""
        if self.recorder_thread and self.recorder_thread.isRunning():
            self.recorder_thread.stop()
            if self.timer_thread and self.timer_thread.isRunning():
                self.timer_thread.stop()
            self.status_label.setText("正在保存...")
            self.status_label.setStyleSheet("color: orange; font-size: 14pt; font-weight: bold;")
            
            # 更新托盘图标
            if self.tray_icon:
                self.tray_icon.setToolTip("mars录屏 - 正在保存")

    def on_recording_stopped(self):
        """录制停止后的处理"""
        self.status_label.setText("录制完成")
        self.status_label.setStyleSheet("color: green; font-size: 14pt; font-weight: bold;")
        self.start_btn.setEnabled(True)
        self.stop_btn.setEnabled(False)
        self.select_path_btn.setEnabled(True)
        self.fps_spin.setEnabled(True)
        self.monitor_combo.setEnabled(True)
        self.audio_checkbox.setEnabled(True)
        
        # 更新托盘图标
        if self.tray_icon:
            self.tray_icon.setToolTip("mars录屏")
            self.tray_record_action.setText("开始录制")
            self.tray_icon.showMessage("录制完成", f"录制已完成并保存至:\n{self.output_file}", QSystemTrayIcon.Information, 3000)
        
        # 显示录制完成消息
        duration_str = self.time_label.text()
        QMessageBox.information(self, "完成", f"录制已完成并保存至:\n{self.output_file}\n\n录制时长: {duration_str}")
        
        # 重置进度条和时间显示
        self.progress_bar.setValue(0)

    def on_recording_error(self, error_msg):
        """处理录制错误"""
        QMessageBox.critical(self, "录制错误", error_msg)
        self.stop_recording()

    def update_recording_time(self, seconds):
        """更新录制时间显示"""
        self.recording_duration = seconds
        hours = seconds // 3600
        minutes = (seconds % 3600) // 60
        secs = seconds % 60
        time_str = f"{hours:02d}:{minutes:02d}:{secs:02d}"
        self.time_label.setText(time_str)

    def update_progress_bar(self, seconds):
        """更新进度条"""
        # 进度条显示录制时长，每小时为一个周期
        progress = (seconds % 3600) / 36  # 0-100
        self.progress_bar.setValue(int(progress))

    def update_preview(self, frame):
        """更新预览画面"""
        try:
            # 调整帧大小以适应预览区域
            h, w, c = frame.shape
            preview_w = self.preview_label.width()
            preview_h = self.preview_label.height()
            
            # 保持宽高比
            aspect_ratio = w / h
            if preview_w / preview_h > aspect_ratio:
                new_w = int(preview_h * aspect_ratio)
                new_h = preview_h
            else:
                new_w = preview_w
                new_h = int(preview_w / aspect_ratio)
            
            # 调整大小
            resized_frame = cv2.resize(frame, (new_w, new_h))
            
            # 转换为QImage并显示
            bytes_per_line = 3 * new_w
            q_img = QImage(resized_frame.data, new_w, new_h, bytes_per_line, QImage.Format_RGB888).rgbSwapped()
            self.preview_label.setPixmap(QPixmap.fromImage(q_img))
        except Exception as e:
            print(f"更新预览失败: {str(e)}")

    def closeEvent(self, event):
        """窗口关闭事件处理"""
        if self.recorder_thread and self.recorder_thread.isRunning():
            reply = QMessageBox.question(
                self, "确认", "录制正在进行中，确定要退出吗？", QMessageBox.Yes | QMessageBox.No
            )
            if reply == QMessageBox.Yes:
                self.recorder_thread.stop()
                if self.timer_thread:
                    self.timer_thread.stop()
                self.recorder_thread.wait()  # 等待线程结束
                event.accept()
            else:
                event.ignore()
        else:
            event.accept()

    def hideEvent(self, event):
        """窗口隐藏事件处理"""
        if self.tray_icon and self.tray_icon.isVisible():
            # 如果有托盘图标，则最小化到托盘
            if self.recorder_thread and self.recorder_thread.isRunning():
                self.tray_icon.showMessage("录制继续中", "应用已最小化到系统托盘，录制继续进行", QSystemTrayIcon.Information, 2000)
            event.accept()
        else:
            event.accept()

if __name__ == "__main__":
    # 设置高DPI支持
    QApplication.setAttribute(Qt.AA_EnableHighDpiScaling, True)
    QApplication.setAttribute(Qt.AA_UseHighDpiPixmaps, True)
    
    app = QApplication(sys.argv)
    app.setStyle("Fusion")  # 使用Fusion风格，在所有平台上看起来更一致
    
    # 设置应用信息
    app.setApplicationName("mars录屏")
    app.setApplicationVersion("1.0.0")
    
    window = ScreenRecorderApp()
    window.show()
    sys.exit(app.exec_())